package com.example.demo.spackTest;

/**
 * 随机森林  回归问题
 * @author ysh   1208706282
 *
 */

import org.apache.log4j.Level;
import org.apache.log4j.Logger;
import org.apache.spark.ml.Pipeline;
import org.apache.spark.ml.PipelineModel;
import org.apache.spark.ml.PipelineStage;
import org.apache.spark.ml.classification.RandomForestClassificationModel;
import org.apache.spark.ml.classification.RandomForestClassifier;
import org.apache.spark.ml.evaluation.MulticlassClassificationEvaluator;
import org.apache.spark.ml.feature.IndexToString;
import org.apache.spark.ml.feature.StringIndexer;
import org.apache.spark.ml.feature.StringIndexerModel;
import org.apache.spark.ml.feature.VectorIndexer;
import org.apache.spark.ml.feature.VectorIndexerModel;
import org.apache.spark.sql.Dataset;
import org.apache.spark.sql.Row;
import org.apache.spark.sql.SparkSession;

public class RandomForest {

    // 使用log4j打印日志
    private static Logger logger = Logger.getLogger(RandomForest.class);

    public static void main(String[] args) {
        SparkSession spark=SparkSession
                .builder()
                .appName("CoFilter")
                .master("local[*]")
                .config("spark.sql.warehouse.dir",
                        "file///:G:/Projects/Java/Spark/spark-warehouse" )
                .getOrCreate();

        /**
         * 屏蔽spark的INFO日志
         */
        logger.getLogger("org.apache.spark").setLevel(Level.ERROR);
        logger.getLogger("org.apache.hadoop").setLevel(Level.ERROR);
        logger.getLogger("org.apache.zookeeper").setLevel(Level.WARN);
        logger.getLogger("org.apache.hive").setLevel(Level.WARN);
        SparkSession.builder().getOrCreate().sparkContext().setLogLevel("WARN");

        // 加载并解析数据文件，然后将其转换为DataFrame。
        Dataset<Row> data = spark.read().format("libsvm").load("D:\\workspace\\Graduation-project\\src\\main\\resources\\python\\CQPython\\static\\question\\TrainDataLibSVM.txt");

        //索引标签，将元数据添加到标签列。
        // 适合整个数据集以将所有标签包括在索引中。
        StringIndexerModel labelIndexer = new StringIndexer()
                .setInputCol("label")
                .setOutputCol("indexedLabel")
                .fit(data);
    //自动识别分类特征并为其编制索引。
    // 设置maxCategories，以便具有> 4个不同值的要素被视为连续要素。
        VectorIndexerModel featureIndexer = new VectorIndexer()
                .setInputCol("features")
                .setOutputCol("indexedFeatures")
                .setMaxCategories(4)
                .fit(data);

        //将数据分为训练集和测试集（保留30％进行测试）
        Dataset<Row>[] splits = data.randomSplit(new double[] {0.7, 0.3});
        Dataset<Row> trainingData = splits[0];
        Dataset<Row> testData = splits[1];

        //训练RandomForest模型。
        RandomForestClassifier rf = new RandomForestClassifier()
                .setLabelCol("indexedLabel")
                .setFeaturesCol("indexedFeatures");

        //将索引标签转换回原始标签。
        IndexToString labelConverter = new IndexToString()
                .setInputCol("prediction")
                .setOutputCol("predictedLabel")
                .setLabels(labelIndexer.labels());

        //管道中的链索引器和林
        Pipeline pipeline = new Pipeline()
                .setStages(new PipelineStage[] {labelIndexer, featureIndexer, rf, labelConverter});

        //训练模型。这也会运行索引器。
        PipelineModel model = pipeline.fit(trainingData);

        // 作出预测。
        Dataset<Row> predictions = model.transform(testData);

        logger.warn("=================================================================================");
        //选择要显示的示例行。
        predictions.select("predictedLabel", "label", "features").show(5);
        logger.warn("=================================================================================");
        //选择（预测，真实标签）并计算测试错误
        MulticlassClassificationEvaluator evaluator = new MulticlassClassificationEvaluator()
                .setLabelCol("indexedLabel")
                .setPredictionCol("prediction")
                .setMetricName("accuracy");
        double accuracy = evaluator.evaluate(predictions);
        System.out.println("Test Error = " + (1.0 - accuracy));

        RandomForestClassificationModel rfModel = (RandomForestClassificationModel)(model.stages()[2]);
        System.out.println("Learned classification forest model:\n" + rfModel.toDebugString());
    }
}

